# -*- coding: utf-8 -*-
from onnx import helper
from onnx import TensorProto
import numpy as np
import random
from .MatMul import MatMul
import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH

#int8类型会报错
#output_types = ['tensor(uint8)', 'tensor(int8)']
output_types = ['tensor(uint8)', 'tensor(uint8)']

class QuantizeLinear():
    def input_shape_correct(self, shapes):
        return True

    def infer_output_shape(self, shapes):
        return shapes[0], True

    def get_initializer_name_list(self):
        initializer_name_list = ['y_scale', 'y_zero_point']
        return initializer_name_list

    def generate_random_initializer(self, in_shapes, out_shape, dtype):
        zp_type = TensorProto.UINT8
        if dtype == 'tensor(int8)': 
            zp_type = TensorProto.INT8

        y_scale_shape = [1]
        y_zp_shape = [1]
        v3 = np.random.randint(0, 100)
        y_scale_array = np.random.random_sample(1)
        y_scale_array = v3*y_scale_array
        y_zp_array = [np.random.randint(0, 256)]
        if dtype == 'tensor(int8)':
            y_zp_array = [np.random.randint(-127, 128+1)]

        '''
        r = random.randint(1, 5)
        if r % 2 == 0:
            a_scale_shape = in_shapes[0].copy()
            a_zp_shape = in_shapes[0].copy()

            size = 1
            for d in a_scale_shape:
                size = size * d
            a_scale_array = np.random.random_sample(size)
            #a_scale_array = float(v1*a_scale_array)

            a_zp_array = np.random.randint(0, 256, size)
            if dtype == 'tensor(int8)':
                a_zp_array = np.random.randint(-127, 128+1, size)

            b_scale_shape = in_shapes[1].copy()
            b_zp_shape = in_shapes[1].copy()

            size = 1
            for d in b_scale_shape:
                size = size * d
            b_scale_array = np.random.random_sample(size)
            #b_scale_array = v2*b_scale_array
            b_zp_array = np.random.randint(0, 256, size)
            if dtype == 'tensor(int8)':
                b_zp_array = np.random.randint(-127, 128+1, size)

            y_scale_shape = out_shape.copy()
            y_zp_shape = out_shape.copy()

            size = 1
            for d in y_scale_shape:
                size = size * d

            y_scale_array = np.random.random_sample(size)
            #y_scale_array = v3*y_scale_array
            y_zp_array = np.random.randint(0, 256, size)
            if dtype == 'tensor(int8)':
                y_zp_array = np.random.randint(-127, 128+1, size)    
        '''

        print('y_scale_shape: ', y_scale_shape)
        print('y_scale_array: ', y_scale_array)
        print('y_zp_shape: ', y_zp_shape)
        print('y_zp_array: ', y_zp_array)

        y_scale = helper.make_tensor('y_scale', TensorProto.FLOAT, y_scale_shape, y_scale_array)    
        y_zero_point = helper.make_tensor('y_zero_point', zp_type, y_zp_shape, y_zp_array)

        initializer = [y_scale, y_zero_point]

        return initializer

    def get_input_name(self):
        input_name = ['x', 'y_scale', 'y_zero_point']
        return input_name

    def get_max_attributes_combination(self):
        return 1

    def get_attributes_by_index_or_random(self, shapes, index):
        #if index % 2 == 0:
        print('+++++++++++++ shapes:', shapes[0])
        return {}, shapes[0]
        '''
        attribute = {}
        attribute['axis'] = random.randint(-1*len(shapes[0]), len(shapes[0])-1)
        '''
        #print('------------ attribute: ', attribute, ', shapes:', shapes)

        return attribute, shapes[0]

    def get_output_type(self, index):
        num = index%2
        print('xxxxx num=', num)
        return output_types[num]


